{
 "cells": [
  {
   "metadata": {},
   "cell_type": "markdown",
   "source": [
    "前面，一个i1只会对应一个q，一个k, 一个v，所以叫单头。\n",
    "\n",
    "而所谓多头，比如两头，就是一个i1会对应两个q，两个k，两个v，所以叫多头。"
   ],
   "id": "fb712db5ee1d4280"
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-07-08T12:52:15.778851Z",
     "start_time": "2025-07-08T12:52:15.774724Z"
    }
   },
   "cell_type": "code",
   "source": [
    "import torch\n",
    "\n",
    "# x表示一个输入的seq\n",
    "x = torch.rand(4, 3)\n",
    "\n",
    "# w是二维矩阵\n",
    "wq_head1 = torch.rand(3, 3)\n",
    "wq_head2 = torch.rand(3, 3)\n",
    "wk_head1 = torch.rand(3, 3)\n",
    "wk_head2 = torch.rand(3, 3)\n",
    "wv_head1 = torch.rand(3, 3)\n",
    "wv_head2 = torch.rand(3, 3)\n",
    "\n",
    "q_head1 = torch.matmul(x, wq_head1)\n",
    "k_head1 = torch.matmul(x, wk_head1)\n",
    "v_head1 = torch.matmul(x, wv_head1)\n",
    "\n",
    "q_head2 = torch.matmul(x, wq_head2)\n",
    "k_head2 = torch.matmul(x, wk_head2)\n",
    "v_head2 = torch.matmul(x, wv_head2)"
   ],
   "id": "7ee18446f5ae7452",
   "outputs": [],
   "execution_count": 48
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-07-08T12:52:17.500905Z",
     "start_time": "2025-07-08T12:52:17.497413Z"
    }
   },
   "cell_type": "code",
   "source": "q_head1",
   "id": "334066ed7e3e9031",
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[0.8723, 0.7370, 0.1824],\n",
       "        [0.5595, 0.8161, 0.2193],\n",
       "        [1.4988, 1.3963, 0.3997],\n",
       "        [0.5766, 0.9081, 0.2250]])"
      ]
     },
     "execution_count": 49,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "execution_count": 49
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-07-08T12:52:24.403739Z",
     "start_time": "2025-07-08T12:52:24.399171Z"
    }
   },
   "cell_type": "code",
   "source": [
    "wq = torch.rand(3, 6)\n",
    "wk = torch.rand(3, 6)\n",
    "wv = torch.rand(3, 6)\n",
    "q = torch.matmul(x, wq)\n",
    "k = torch.matmul(x, wk)\n",
    "v = torch.matmul(x, wv)\n",
    "q"
   ],
   "id": "dcdc05ea7de5e68d",
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[1.0413, 0.4550, 0.1371, 0.9614, 0.7707, 0.6358],\n",
       "        [0.3589, 0.6196, 0.3238, 0.1411, 0.3041, 0.4545],\n",
       "        [1.0109, 0.9428, 0.4445, 0.9574, 0.6860, 0.9082],\n",
       "        [0.6025, 0.6861, 0.3268, 0.2804, 0.5271, 0.5670]])"
      ]
     },
     "execution_count": 50,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "execution_count": 50
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-07-08T12:53:07.834937Z",
     "start_time": "2025-07-08T12:53:07.831926Z"
    }
   },
   "cell_type": "code",
   "source": [
    "q_head1, q_head2 = torch.split(q, 3, dim=1)\n",
    "k_head1, k_head2 = torch.split(k, 3, dim=1)\n",
    "v_head1, v_head2 = torch.split(v, 3, dim=1)"
   ],
   "id": "f30456f6840caa46",
   "outputs": [],
   "execution_count": 52
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-07-08T12:53:14.673274Z",
     "start_time": "2025-07-08T12:53:14.670146Z"
    }
   },
   "cell_type": "code",
   "source": "q_head1",
   "id": "b5c369406902a0fc",
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[1.0413, 0.4550, 0.1371],\n",
       "        [0.3589, 0.6196, 0.3238],\n",
       "        [1.0109, 0.9428, 0.4445],\n",
       "        [0.6025, 0.6861, 0.3268]])"
      ]
     },
     "execution_count": 53,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "execution_count": 53
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-07-08T12:47:03.690631Z",
     "start_time": "2025-07-08T12:47:03.687514Z"
    }
   },
   "cell_type": "code",
   "source": [
    "qk_score_head1 = torch.matmul(q_head1, k_head1.T)\n",
    "dk_head1 = k_head1.size(-1)\n",
    "qk_weight_head1 = torch.softmax(qk_score_head1 / torch.sqrt(torch.tensor(dk_head1)), dim=-1)\n",
    "o_head1 = torch.matmul(qk_weight_head1, v_head1)"
   ],
   "id": "7b7d504b707c1b3",
   "outputs": [],
   "execution_count": 42
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-07-08T12:47:09.357859Z",
     "start_time": "2025-07-08T12:47:09.353436Z"
    }
   },
   "cell_type": "code",
   "source": [
    "qk_score_head2 = torch.matmul(q_head2, k_head2.T)\n",
    "dk_head2 = k_head2.size(-1)\n",
    "qk_weight_head2 = torch.softmax(qk_score_head2 / torch.sqrt(torch.tensor(dk_head2)), dim=-1)\n",
    "o_head2 = torch.matmul(qk_weight_head2, v_head2)\n",
    "\n",
    "o_head1, o_head2"
   ],
   "id": "7eebecc244dbf577",
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(tensor([[1.2538, 0.8411, 1.7767],\n",
       "         [1.3108, 0.8815, 1.8632],\n",
       "         [1.3340, 0.8973, 1.8966],\n",
       "         [1.3448, 0.9051, 1.9136]]),\n",
       " tensor([[1.5382, 0.8595, 0.6576],\n",
       "         [1.6358, 0.9160, 0.6935],\n",
       "         [1.6440, 0.9181, 0.6949],\n",
       "         [1.6767, 0.9330, 0.7049]]))"
      ]
     },
     "execution_count": 43,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "execution_count": 43
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-07-08T12:47:53.606896Z",
     "start_time": "2025-07-08T12:47:53.602823Z"
    }
   },
   "cell_type": "code",
   "source": [
    "o = torch.cat([o_head1, o_head2], dim=1)\n",
    "o"
   ],
   "id": "38e503043ee37c12",
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[1.2538, 0.8411, 1.7767, 1.5382, 0.8595, 0.6576],\n",
       "        [1.3108, 0.8815, 1.8632, 1.6358, 0.9160, 0.6935],\n",
       "        [1.3340, 0.8973, 1.8966, 1.6440, 0.9181, 0.6949],\n",
       "        [1.3448, 0.9051, 1.9136, 1.6767, 0.9330, 0.7049]])"
      ]
     },
     "execution_count": 44,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "execution_count": 44
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-07-08T12:49:12.394285Z",
     "start_time": "2025-07-08T12:49:12.390360Z"
    }
   },
   "cell_type": "code",
   "source": [
    "wo = torch.rand(6, 10)\n",
    "out_put = torch.matmul(o, wo)\n",
    "out_put"
   ],
   "id": "18cc2965ef964b2",
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[2.8993, 1.8304, 1.9606, 1.8212, 3.5667, 2.8518, 2.8931, 2.9325, 3.0127,\n",
       "         2.4597],\n",
       "        [3.0667, 1.9280, 2.0576, 1.9154, 3.7609, 3.0047, 3.0429, 3.0969, 3.1720,\n",
       "         2.5826],\n",
       "        [3.0947, 1.9535, 2.0881, 1.9408, 3.7995, 3.0401, 3.0837, 3.1308, 3.2111,\n",
       "         2.6195],\n",
       "        [3.1394, 1.9764, 2.1096, 1.9623, 3.8496, 3.0750, 3.1192, 3.1762, 3.2474,\n",
       "         2.6459]])"
      ]
     },
     "execution_count": 46,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "execution_count": 46
  },
  {
   "metadata": {},
   "cell_type": "markdown",
   "source": "多头注意力就是增加了另外的q、k、v，从而可以从多个角度来捕捉token之间的相关性。",
   "id": "4605d9dfe0d35435"
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-07-08T13:05:30.809511Z",
     "start_time": "2025-07-08T13:05:30.802038Z"
    }
   },
   "cell_type": "code",
   "source": [
    "import torch\n",
    "import torch.nn as nn\n",
    "\n",
    "\n",
    "class MultiHeadAttention(nn.Module):\n",
    "\n",
    "    def __init__(self, embed_dim: int, attn_dim: int, output_dim: int, num_heads: int):\n",
    "        super().__init__()\n",
    "\n",
    "        self.embed_dim = embed_dim\n",
    "        self.attn_dim = attn_dim\n",
    "        self.output_dim = output_dim\n",
    "        self.num_heads = num_heads\n",
    "        self.head_dim = attn_dim // num_heads # //表示向下取整，attn_dim是head_dim的整数倍\n",
    "\n",
    "        # QKV投影层：从输入维度映射到内部维度\n",
    "        # projection\n",
    "        self.q_proj = nn.Linear(embed_dim, self.attn_dim)\n",
    "        self.k_proj = nn.Linear(embed_dim, self.attn_dim)\n",
    "        self.v_proj = nn.Linear(embed_dim, self.attn_dim)\n",
    "\n",
    "        # 输出投影层：从内部维度映射到输出维度\n",
    "        self.out_proj = nn.Linear(self.attn_dim, self.output_dim)\n",
    "\n",
    "    def forward(self, x):\n",
    "        \"\"\"\n",
    "        输入: [batch_size, seq_len, embed_dim]\n",
    "        返回: [batch_size, seq_len, output_dim]\n",
    "        \"\"\"\n",
    "        batch_size, seq_len, embed_dim = x.shape\n",
    "\n",
    "        # 投影到QKV空间\n",
    "        q = self.q_proj(x)  # [batch_size, seq_len, attn_dim]\n",
    "        k = self.k_proj(x)  # [batch_size, seq_len, attn_dim]\n",
    "        v = self.v_proj(x)  # [batch_size, seq_len, attn_dim]\n",
    "\n",
    "        # [batch_size, seq_len, num_heads, head_dim]\n",
    "        # 分割多头 [batch_size, num_heads, seq_len, head_dim]\n",
    "        q = q.view(batch_size, seq_len, self.num_heads, self.head_dim).transpose(1, 2)\n",
    "        k = k.view(batch_size, seq_len, self.num_heads, self.head_dim).transpose(1, 2)\n",
    "        v = v.view(batch_size, seq_len, self.num_heads, self.head_dim).transpose(1, 2)\n",
    "\n",
    "        # 计算注意力得分\n",
    "        # q   [batch_size, num_heads, seq_len, head_dim]\n",
    "        # k.T [batch_size, num_heads, head_dim, seq_len]\n",
    "        # q @ k.T 形状: [batch_size, num_heads, seq_len, seq_len]\n",
    "        attn_scores = torch.matmul(q, k.transpose(-2, -1))\n",
    "\n",
    "        # 缩放因子：防止乘积过大\n",
    "        d_k = k.size(-1)\n",
    "        attn_scores = attn_scores / torch.sqrt(torch.tensor(d_k))\n",
    "\n",
    "        # 计算注意力权重\n",
    "        attn_weights = torch.softmax(attn_scores, dim=-1)\n",
    "\n",
    "        # 计算注意力输出\n",
    "        # attn_weights [batch_size, num_heads, seq_len, seq_len]\n",
    "        # v            [batch_size, num_heads, seq_len, head_dim]\n",
    "        # [batch_size, num_heads, seq_len, head_dim]\n",
    "        attn_out = torch.matmul(attn_weights, v)\n",
    "\n",
    "        # 合并多头 [batch_size, seq_len, attn_dim]\n",
    "\n",
    "        # [batch_size, seq_len, num_heads, head_dim]\n",
    "        attn_out = attn_out.transpose(1, 2).reshape(batch_size, seq_len, self.attn_dim)\n",
    "\n",
    "        # 投影到输出空间\n",
    "        return self.out_proj(attn_out)"
   ],
   "id": "403c5a8fcae45aae",
   "outputs": [],
   "execution_count": 54
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-07-08T13:06:46.567770Z",
     "start_time": "2025-07-08T13:06:46.563512Z"
    }
   },
   "cell_type": "code",
   "source": [
    "x = torch.rand(1, 4, 2)\n",
    "\n",
    "attn = MultiHeadAttention(embed_dim=2, attn_dim=6, output_dim=8, num_heads=2)\n",
    "out = attn(x)\n",
    "\n",
    "print(x.shape)\n",
    "print(out.shape)"
   ],
   "id": "6c229fa626c5ff1c",
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "torch.Size([1, 4, 2])\n",
      "torch.Size([1, 4, 8])\n"
     ]
    }
   ],
   "execution_count": 62
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 2
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython2",
   "version": "2.7.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
